import torch.nn as nn
import torch

"""
    逻辑斯蒂回归（分类问题）
"""
# 1、构造数据集
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[0], [0], [1]])


# 2、设计模型
class LogisticRegressionModel(nn.Module):
    def __init__(self):
        super(LogisticRegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        y_pre = torch.sigmoid(self.linear(x))  # 将线性预测值进行sigmoid函数激活
        return y_pre


model = LogisticRegressionModel()

# 3、构造损失函数和优化器
"""
BCELoss(self, weight: Optional[Tensor] = None, size_average=None, 
reduce=None, reduction: str = 'mean')
reduction默认平均
"""
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

if __name__ == "__main__":
    # 4、训练模型
    # 前馈，反馈，更新模型
    for epoch in range(10000):
        y_pre = model(x_data)
        loss = criterion(y_pre, y_data)
        print("第" + str(epoch) + "次迭代，loss=" + str(loss.item()))

        optimizer.zero_grad()  # 将所有权重的梯度归零
        loss.backward()  # 反向传播
        optimizer.step()  # 权重更新

    print("w= ", model.linear.weight.item())
    print("b= ", model.linear.bias.item())
    x_test = torch.Tensor([[4.0]])
    y_test = model(x_test)
    print("y_pred = ", y_test.data)
